import torch
import torch.nn.functional as F
import rl_utils
import copy



class PolicyNetContinuous(torch.nn.Module):
    # 正切函数的值域是[-1,1]
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNetContinuous, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)
        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        mu = 2.0 * torch.tanh(self.fc_mu(x))
        std = F.softplus(self.fc_std(x))
        return mu, std  # 高斯分布的均值和标准差


class ValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super(ValueNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)


class TRPOContinuous:
    """ 处理连续动作的TRPO算法 """
    def __init__(self, hidden_dim, state_space, action_space, lmbda, kl_constraint, alpha, critic_lr, discount_factor, device):
        state_dim = state_space.shape[0]
        action_dim = action_space.shape[0]
        
        # 策略网络参数不需要优化器更新
        self.actor = PolicyNetContinuous(state_dim, hidden_dim, action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),  lr=critic_lr)
        
        self.discount_factor = discount_factor
        self.lmbda = lmbda  # 广义优势估计参数
        self.kl_constraint = kl_constraint  # KL距离最大限制
        self.alpha = alpha  # 线性搜索参数
        self.device = device

    def take_action(self, state):
        state = torch.tensor(state, dtype=torch.float).to(self.device)
        mu, std = self.actor(state)
        action_dist = torch.distributions.Normal(mu, std)
        action = action_dist.sample()
        return [action.item()]

    def hessian_matrix_vector_product(self,states,old_action_dists,vector,damping=0.1):
        # 计算黑塞矩阵和一个向量的乘积
        mu, std = self.actor(states)
        new_action_dists = torch.distributions.Normal(mu, std)
        kl = torch.mean(torch.distributions.kl.kl_divergence(old_action_dists,new_action_dists))  # 计算平均KL距离
        kl_grad = torch.autograd.grad(kl,self.actor.parameters(),create_graph=True)
        kl_grad_vector = torch.cat([grad.view(-1) for grad in kl_grad])
        # KL距离的梯度和向量进行点积运算
        kl_grad_vector_product = torch.dot(kl_grad_vector, vector)
        grad2 = torch.autograd.grad(kl_grad_vector_product,self.actor.parameters())
        grad2_vector = torch.cat([grad.contiguous().view(-1) for grad in grad2])
        return grad2_vector + damping * vector

    def conjugate_gradient(self, grad, states, old_action_dists):
        # 共轭梯度法求解方程
        x = torch.zeros_like(grad) # x 是全零向量，维度和 g 一样
        r = grad.clone() # r = g - H * x = g
        p = grad.clone() # p = r
        rdotr = torch.dot(r, r)
        for i in range(10): # 共轭梯度主循环
            Hp = self.hessian_matrix_vector_product(states, old_action_dists,p)
            alpha = rdotr / torch.dot(p, Hp)
            x += alpha * p
            r -= alpha * Hp
            new_rdotr = torch.dot(r, r)
            if new_rdotr < 1e-10:
                break
            beta = new_rdotr / rdotr
            p = r + beta * p
            rdotr = new_rdotr
        return x

    def compute_surrogate_obj(self, states, actions, advantage, old_log_probs,actor): # 计算策略目标
        mu, std = actor(states)
        action_dists = torch.distributions.Normal(mu, std)
        log_probs = action_dists.log_prob(actions)
        ratio = torch.exp(log_probs - old_log_probs)
        return torch.mean(ratio * advantage)

    def line_search(self, states, actions, advantage, old_log_probs,old_action_dists, max_vec):
        # 用线性搜索找到一个最小非负整数i，其能提升策略并满足 KL 距离限制
        old_para = torch.nn.utils.convert_parameters.parameters_to_vector(self.actor.parameters())
        old_obj = self.compute_surrogate_obj(states, actions, advantage,old_log_probs, self.actor)
        for i in range(15):
            # alpha 是一个决定线性搜索长度的超参数
            # i 是非负整数
            coef = self.alpha**i
            # 新的策略参数
            new_para = old_para + coef * max_vec
            # 新的策略
            new_actor = copy.deepcopy(self.actor)
            torch.nn.utils.convert_parameters.vector_to_parameters(new_para, new_actor.parameters())
            mu, std = new_actor(states)
            new_action_dists = torch.distributions.Normal(mu, std)
            # KL散度，衡量新旧策略之间的距离
            kl_div = torch.mean(torch.distributions.kl.kl_divergence(old_action_dists,new_action_dists))
            # 替代优化目标
            new_obj = self.compute_surrogate_obj(states, actions, advantage,old_log_probs, new_actor)
            # 如果新的策略参数可以提升策略并满足 KL 距离限制
            if new_obj > old_obj and kl_div < self.kl_constraint:
                return new_para
        return old_para

    def policy_learn(self, states, actions, old_action_dists, old_log_probs,advantage):
        # 计算策略目标
        surrogate_obj = self.compute_surrogate_obj(states, actions, advantage,old_log_probs, self.actor)
        # 计算策略目标函数的梯度 g
        grads = torch.autograd.grad(surrogate_obj, self.actor.parameters())
        obj_grad = torch.cat([grad.view(-1) for grad in grads]).detach()
        # 用共轭梯度法计算 x = H^(-1) * g，即解 H * x = g
        descent_direction = self.conjugate_gradient(obj_grad, states,old_action_dists)
        # 计算 H * x
        Hd = self.hessian_matrix_vector_product(states, old_action_dists,descent_direction)
        # 计算线性搜索中的根号
        max_coef = torch.sqrt(2 * self.kl_constraint /(torch.dot(descent_direction, Hd) + 1e-8))
        # 策略参数更新的最大步长 = descent_direction * max_coef
        # 线性搜索得到新的策略参数
        new_para = self.line_search(states, actions, advantage, old_log_probs,old_action_dists,descent_direction * max_coef)
        # 用线性搜索后的参数更新策略
        torch.nn.utils.convert_parameters.vector_to_parameters(new_para, self.actor.parameters())

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions'],dtype=torch.float).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],dtype=torch.float).view(-1, 1).to(self.device)
        rewards = (rewards + 8.0) / 8.0  # 对奖励进行修改,方便训练
        
        # critic 是一个已经学习的状态价值函数
        td_target = rewards + self.discount_factor * self.critic(next_states) * (1 - dones)
        # 时序差分误差
        td_delta = td_target - self.critic(states)
        # 估计每个状态动作对的优势，根据收集到的数据和价值网络
        advantage = rl_utils.compute_advantage(self.discount_factor, self.lmbda,td_delta.cpu()).to(self.device)
        mu, std = self.actor(states)
        old_action_dists = torch.distributions.Normal(mu.detach(),std.detach())
        old_log_probs = old_action_dists.log_prob(actions)
        
        # 更新价值网络 critic
        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        
        # 更新策略网络 actor
        self.policy_learn(states, actions, old_action_dists, old_log_probs,advantage)
        
        

